import asyncio
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from json_repair import repair_json
from typing import Optional, List
from loguru import logger
from template_parser import TemplateParser
from pathlib import Path
import json
import os

os.environ["OPENAI_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"
os.environ["OPENAI_API_KEY"] = "sk-e1ffe36667f84c57bebd8f6f7b120f3a"


class PlanAgent(object):

    def __init__(self):
        self.client = MultiServerMCPClient(
            {
                "sequential-thinking": {
                    "url": "http://127.0.0.1:8000/sse",
                    "transport": "sse",
                }
            }
        )
        self.llm = ChatOpenAI(model="qwen-plus-latest")
        with open(Path(__file__).parent / "format_plan.md", "r", encoding="utf-8") as fs:
            self.plan_template = TemplateParser(fs.read())

    async def create_plan(self, task_description: str, agents_description: List = None) -> str:
        system_prompt = """You are a helpful assistant that can break down complex problems.
            Use the available tools when relevant to perform step-by-step analysis."""
        if agents_description is not None:
           system_prompt += (
                f"\nNow we have {len(agents_description)} agents. "
                f"The infomation of them are below: {json.dumps(agents_description)}\n"
                "When creating steps in the planning tool, please specify the agent names using the format '[agent_name]'."
            )

        tools = await self.client.get_tools()
        memory = MemorySaver()
        agent = create_react_agent(
            model=self.llm,
            tools=tools,
            prompt=system_prompt,
            checkpointer=memory,
            debug=False
        )
        config = {"configurable": {"thread_id": "226"}}
        response = await agent.ainvoke(
            input={"messages": task_description},
            config=config
        )
        messages = response['messages']
        return messages[-1].content

    async def format_plan(
            self, task_description: str, detailed_plan: str,
            agents_description: Optional[List] = None
    ) -> List:
        """
            Forms plan into a JSON format.

            Parameters:
                task_description (str): the task description.
                agents_description (str): available tools description.
                detailed_plan (str): The detailed content of plan for the task.

            Returns:
                list: List of sub-plans with following format:
                   -title: Title for the sub-plan.
                   -steps: List of sub-plan steps.
                   -tools: List of available tools used in the sub-plan.
        """
        prompt = self.plan_template.format(
            inputs={
                "task_description": task_description,
                "agents_description": agents_description,
                "detailed_plan": detailed_plan
            },
            remove_template_variables=True
        )
        output = self.llm.invoke(prompt)
        content = output.content
        think_end_pos = content.find('</think>')
        if think_end_pos != -1:
            content = content[think_end_pos + len('</think>'):]
        result = repair_json(content, return_objects=True)

        # 严格的格式检查。
        if not isinstance(result, list):
            raise Exception(f"输出结果不是一个有效的列表格式：{result}")
        fields = {"name", "steps", "tools"}
        for item in result:
            if not isinstance(item, dict):
                raise Exception(f"结果列表中的子项不是一个有效的字典格式：{item}")
            not_included = fields.difference(item.keys())
            if len(not_included) != 0:
                raise Exception(f"字段{not_included}未包含在子项中：{item}")
        return result

    async def invoke(self, task_description: str, agents_description: Optional[List] = None) -> List:
        str_plan = await self.create_plan(task_description, agents_description)
        logger.info(str_plan)
        try_time = 3
        while try_time != 0:
            try:
                json_plan = await self.format_plan(task_description, str_plan, agents_description)
                logger.info(json_plan)
                return json_plan
            except Exception as e:
                logger.error(e)
                try_time -= 1
        return []


if __name__ == '__main__':
    agent = PlanAgent()
    final_plan = asyncio.run(agent.invoke("将`泡一杯茶`分拆为多个可执行的步骤"))
    print(json.dumps(final_plan, ensure_ascii=False, indent=4))
